/**
 * Friendly Interactive Recursion Educator (FIRE)
 * COMS 4115: Programming Languages & Translators
 * Prof. Aho
 * 
 * @author Brian Lu
 * @author Sahar Hasan
 * @author Steve Hehl
 * @author Tristan Naumann
 * @version 2009-05-12
 */
grammar Fire;

options {
  language = Java;
  output=template;
  backtrack = true;
}

tokens {
  INDENT;
  DEDENT;
}

@header{ 
package fire;
import java.util.LinkedList;


}

@lexer::header{
package fire;
}

@lexer::members {

public static final int TAB_WIDTH = 4;
}

@parser::members{ 

  public Scope globalScope = new Scope(null);
  public Scope currentScope = globalScope; // "stack top" pointer
  public Map<String,Integer> userRefs = new HashMap<String, Integer>(); 
  

  public ArrayList<String> watchList = new ArrayList<String>(); 
  
  //public ArrayList<String> baseConds = new ArrayList<String>();
  //public ArrayList<String> recurseConds = new ArrayList<String>();
  
  public ArrayList<String> condList = new ArrayList<String>();
  
  public ArrayList<String> paramList = new ArrayList<String>(); 
  public StandardFunctions stdfunctions = new StandardFunctions();
  
  public boolean inRecursiveFunc = false; 
  
  public FunctionSymbol currFunction = null;
  
  public void doError(String msg) {
    System.out.println("ERROR: "+ msg);
    System.exit(0);
  }
  
  public void doError(int line, String msg) {
    System.out.println("ERROR @"+ line +": "+ msg);
    System.exit(0);
  }
  
  public Symbol resolveType(String name, Symbol ... args) {

    /** get list of matching signatures for name **/
    LinkedList<ArrayList<Symbol>> signatures = stdfunctions.lookupSignatures(name);
    if (signatures == null) doError("Unknown operator "+ name +".");
    
    /** look for matching signature **/
    Symbol r = null;
    for (int j = 0; j < signatures.size(); j++) {
      ArrayList<Symbol> sig = signatures.get(j);
      
      if (sig.size() != args.length + 1) continue;  // not a match
      
      /** try to match **/
      boolean match = true;
      for (int i = 0; i < args.length; i++) {
        if (!args[i].sameType(sig.get(i+1))) {
          match = false;
          break;
        } 
        
      }
      
      if (match) return sig.get(0); // match found
    }
    return r;  // could not find match 
  }
  
  public Symbol handleFunctionCall(String name, int line, ArrayList<Symbol> args){
     
  // if(builtFunctions.contains(name))
    //  System.out.println("Found built-in"); 
      //Return StringTemplate or String for built-in 

    FunctionSymbol type = (FunctionSymbol)currentScope.lookup(name); 
   
    if (type == null){ //Looking out for forward references
      System.out.println("ERROR @"+ line +": function '"+ name +"' not defined.");
      return null;
    } //Return StringTemplate for actual function call 
    else{
      ArrayList<Symbol> sig = type.getSignature();
      /** try to match **/
      boolean match = true;
      for(int i = 0; i < args.size(); i++){
        if (!sig.get(i+1).sameType(args.get(i))) {
          if(resolveType("assign", sig.get(i+1), args.get(i)) == null){
            match = false;
            break;
          }
        } 
      }
      if (match) return type;
    }
    return null;  
    
  }
  
  public void handleForwardRefs(){ 

    java.util.Set<String> refs = userRefs.keySet();
    java.util.Iterator iter = refs.iterator(); 
    while(iter.hasNext()){
      String key = (String) iter.next();
      if(currentScope.lookup(key) == null){
        int line = userRefs.get(key).intValue();
        doError(line, "function '"+ key + "' not defined.");
      }   
    }
  }
  
  
  public void defaultWatchList(StringTemplate t){
  
 
    ArrayList temp = (ArrayList)t.getAttribute("list");

    //No formal parameters--ridiculous case...
    if(temp == null){ 
      watchList.add(null);
      return;
    }
      
    String def; 
    for(int i = 0; i < temp.size(); i++){ 
      def = temp.get(i).toString(); 
      watchList.add(def.substring(def.indexOf(" ")+1));
    }
     
  }
  
  public void setWatchList(List l){
  
    watchList.clear();
    
    if(l==null)
      return;
    
    String s;   
    for(int i = 0; i < l.size(); i++){
      s = ((CommonToken)l.get(i)).getText(); 
      watchList.add(s); 
     }
     
  }
  
  public String buildKleene(List ... lists){ 
 
      String output = new String();

      if(lists[0] == null)
        return output;
        
      int totalStars = lists[0].size();
      int totalLists = lists.length; 
        
      for(int i = 0; i < totalStars; i++){
        for (int j = 0; j < totalLists; j++){
          if(lists[j] != null && lists[j].get(i) != null){      
            Class type = lists[j].get(i).getClass();
            if (type.getName().endsWith("CommonToken"))
              output += ( (CommonToken)(lists[j].get(i)) ).getText();
            else if(type.getName().endsWith("StringTemplate"))
              output += ( (StringTemplate)(lists[j].get(i)) ).toString();
            else{
              output += lists[j].get(i).toString();
              }
          }
        }
      }
    
      return output;
  
    }
  
    public String buildKleeneTests(List ... lists){ 
 
      
      String output = new String();

     if(lists[0] == null)
        return output;
         
      int totalStars = lists[0].size();
      int totalLists = lists.length; 
    
    
      for(int i = 0; i < totalStars; i++){
        for (int j = 0; j < totalLists; j++){
          Class type = lists[j].get(0).getClass();
          if (type.getName().endsWith("CommonToken")){
            String op = ((CommonToken)(lists[j].get(i))).getText();
            if(op.equals("and"))
              op = "&&";
            else
              op = "||";
            output += " " + op;
          }
          else
            output += " " + ( (StringTemplate)(lists[j].get(i)) ).toString();
        }
      }
      return output;
    }
    
    public ArrayList<String> handleDeclr(ArrayList<String> soFar, List l1, List l2, int num_kleenes, int num_ops, boolean op){
    
      if(op==true){
        soFar.add(num_kleenes,(" " + ((CommonToken)l1.get(num_ops)).getText() + " " + l2.get(num_ops)));
       }
       else{
        soFar.add(num_kleenes, null);
       }
       
       return soFar;
     }
}

    
s : g=program -> {$g.st}
  ;
  
program
@after{
handleForwardRefs();
}
        : (d+=declr_stmt NEWLINE)* (f+=func)* recursive_func
          -> prog(globals={$d},rfunction={$recursive_func.st},function={$f}, watch={watchList}, conds={condList}, params={paramList})
        ;


recursive_func
@init{
  FunctionSymbol funcsymbol = null;
  inRecursiveFunc = true;
  
}
               : DEFINE RECURSIVE t=type n=NAME 
                {funcsymbol = (FunctionSymbol) currentScope.defineFunction($n.text, t.type, true);} 
                {currFunction = funcsymbol;}
                {currentScope = new Scope(currentScope);}
                LPAREN p=parameters RPAREN {defaultWatchList($p.st);} 
                {funcsymbol.setArgs($p.args);} 
                COLON NEWLINE INDENT b=special_block DEDENT 
                {currentScope = currentScope.pop(); inRecursiveFunc = false;}
                -> recursive_function(rtype={$t.st}, name={$n.text}, params={$p.st}, body={$b.st})
              ;
              
special_block : b=base_block r=recursive_block (w=watch_block)? 
                -> special_body(base={$b.st}, recursive={$r.st}, watch={$w.st})
              ; 
              
base_block
            : BASE COLON NEWLINE INDENT (b+=stmt)+ DEDENT
             -> base_body(body={$b})
            ;

recursive_block
                : RECURSIVE COLON NEWLINE INDENT (b+=stmt)+ DEDENT 
                 -> recursive_body(body={$b})
                ;

watch_block : WATCH COLON n+=NAME (COMMA n+=NAME)* {setWatchList($n);};

             
blocks  : COLON NEWLINE 
          {currentScope = new Scope(currentScope);}
          INDENT (b+=stmt)+ DEDENT
          {currentScope = currentScope.pop();}
          -> regular_body(body={$b})
        ;    
        
func
@init{
  FunctionSymbol funcsymbol = null;
 }    
       : DEFINE t=type n=NAME
       {funcsymbol = (FunctionSymbol) currentScope.defineFunction($n.text,t.type);} 
       {currFunction = funcsymbol;}
       {currentScope = new Scope(currentScope);}
       LPAREN p=parameters RPAREN {funcsymbol.setArgs($p.args);} COLON NEWLINE INDENT (b+=stmt)+ DEDENT            
       {currentScope = currentScope.pop();} 
      -> function(rtype={$type.st},name={$NAME.text}, params={$parameters.st}, block={$b})
      ;

stmt    : simple_stmt -> {$simple_stmt.st}
        | compound_stmt -> {$compound_stmt.st}
        | r=cond_stmt NEWLINE {if(inRecursiveFunc){ condList.add($r.text); }}-> {$cond_stmt.st}
        ;


compound_stmt : if_stmt -> {$if_stmt.st}
              | while_stmt -> {$while_stmt.st}
              | for_stmt -> {$for_stmt.st}
              ;
              
if_stmt : IF test b1=blocks (el+=elseif_part)* (ELSE b2=blocks)?
          -> if_statement(cond={$test.st}, block={$b1.st}, elif={$el},else={$ELSE.text}, block2={$b2.st});
 
elseif_part : ELSEIF test blocks -> elseif_statement(cond={$test.st}, block={$blocks.st})
            ;
            
while_stmt : WHILE test blocks -> while_statement(cond={$test.st}, block={$blocks.st})
           ;            

for_stmt : FOR t=type v=NAME 
           {currentScope.defineVariable($v.text, t.type);} 
           IN (target=NAME) blocks 
            -> for_statement(type={$type.st}, var={$v.text}, targ={$target.text}, body={$blocks.st})
         ;
 
simple_stmt : sub_stmt NEWLINE -> simple_statement(stmt={$sub_stmt.st}); 

cond_stmt returns [Symbol type]
          : LCURLY (t=test)? RCURLY x=sub_stmt_prime {Symbol rtype = resolveType("assign", currFunction.type, $x.type); if(rtype==null){doError("return types mismatched.");}} 
          -> {inRecursiveFunc}? gui_cond(cond={$t.st},stat={$x.st}, vars={currentScope.getScopeVariables()})
          -> cond_statement(cond={$t.st}, stat={$x.st})
          ;

sub_stmt_prime returns [Symbol type]  
                : e=expr_stmt {$type = $e.type;} -> {$e.st}
                | t=test {$type = $t.type;} -> {$t.st}
                ;
                
sub_stmt  : declr_stmt -> {$declr_stmt.st}
          | print_stmt -> {$print_stmt.st}
          | expr_stmt -> {$expr_stmt.st}
          ;
            
declr_stmt
@init{
  int num_ops = -1, num_kleenes=-1;
  ArrayList<String> comb = new ArrayList<String>();
  boolean op = false;
  ArrayList<Token> n = new ArrayList<Token>();
  
}
  : t=type NAME {currentScope.defineVariable($NAME.text, t.type);} (ASSIGN t1=test)? 
   (o+=COMMA v=NAME {num_kleenes++; $v.setText("fire_"+$v.getText()); n.add($v);} {currentScope.defineVariable($v.text, t.type);}
   (a+=ASSIGN rhs+=test {num_ops++; op=true;})?
   {comb = handleDeclr(comb, $a, $rhs, num_kleenes, num_ops, op); op=false;} )* 
    -> declr_statement(type={$type.st}, var={$NAME.text}, assign={$t1.st}, more={buildKleene($o,n,comb)})
  ;
        
print_stmt  : PRINT LPAREN (printlist)? RPAREN -> print_statement(out={$printlist.st})
            ;
        
printlist : t+=test (COMMA t+=test)* 
            -> template(list={$t}) "<list; separator=\"+\">"
          ;
          

expr_stmt returns [Symbol type]
          : l=atom a=augassign r=test {Symbol rtype = resolveType("assign", $l.type, $r.type); $type = rtype; if(rtype==null){doError("assignment types mismatched.");}}
            -> expr_statement(lhs={$l.st}, assign={$a.st}, rhs={$r.st}) 
          ;
          
augassign : ASSIGN -> assign(x={$ASSIGN.text})
          ;
                    
test returns [Symbol type]  
      : t=or_test {$type=$t.type;}-> {$t.st}
      ;

or_test returns [Symbol type] 
@init{ 
  $type = null;
  ArrayList<StringTemplate> t2 = new ArrayList<StringTemplate>();
 }  
        : t1=and_test {$type = $t1.type;} (o+=OR t=and_test {$type=resolveType("bop", $type, $t.type); t2.add($t.st);})*
          -> test_or(t={$t1.st}, more={buildKleeneTests($o,t2)})
        ;

and_test returns [Symbol type] 
@init{ 
  $type = null;
  ArrayList<StringTemplate> t2 = new ArrayList<StringTemplate>();
 }  
         : t1=not_test {$type = $t1.type;} (o+=AND t=not_test {$type=resolveType("bop", $type, $t.type); t2.add($t.st);})* 
          -> test_and(t={$t1.st}, more={buildKleeneTests($o,t2)})
         ;
         
not_test returns [Symbol type] 
         : NOT n=not_test {$type = $n.type;}-> test_not(t={$n.st})
         | c=comparison {$type =$c.type;}-> {$c.st}
         ;
         
comparison returns [Symbol type]
@init{
  $type = null;
  ArrayList<StringTemplate> e2 = new ArrayList<StringTemplate>();
}
          : e1=expr {$type = $e1.type;} (o+=comp_op e=expr {String op = $comp_op.st.toString(); if(op.equals("==") || op.equals("!=")) {$type = resolveType("compareeq", $type, $e1.type);} else{ $type=resolveType("compare", $type, $e1.type);} e2.add($e.st);})*
            -> compare(operand={$e1.st}, more={buildKleene($o, e2)})
          ;
          
comp_op : LT -> compare_op(operator={$LT.text})
        | GT -> compare_op(operator={$GT.text})
        | EQ -> compare_op(operator={$EQ.text})
        | GEQ -> compare_op(operator={$GEQ.text})
        | LEQ -> compare_op(operator={$LEQ.text})
        | NEQ -> compare_op(operator={$NEQ.text})
        ;
        
expr returns [Symbol type] : e=arith_expr {$type = $e.type;} -> {$e.st}         
     ;  
     
arith_expr returns [Symbol type]
@init{
  $type = null;
  ArrayList<StringTemplate> e2 = new ArrayList<StringTemplate>();
}
          : e1=term {$type = $e1.type;}((o+=PLUS | o+=MINUS) e=term {$type=resolveType("aop", $type, $e.type); e2.add($e.st);})*
            -> arithexpr(operand={$e1.st},more={buildKleene($o, e2)})          
          ;
          
term returns [Symbol type]
@init{
  $type = null; 
  int curr = 0;
  ArrayList<StringTemplate> f2 = new ArrayList<StringTemplate>();
  
 }   
     : f1=factor {$type = $f1.type;} ((o+=STAR | o+=SLASH | o+=PERCENT) f=factor {$type=resolveType("aop", $type, $f.type); f2.add($f.st);})* 
       -> term(operand={$f1.st}, more={buildKleene($o, f2)})
     ;

factor returns [Symbol type] 
      : PLUS f=factor {$type = resolveType("uop", $f.type);} -> factor(uop={$PLUS.text},operand={$f.st})
       | MINUS f=factor {$type = resolveType("uop", $f.type);}-> factor(uop={$MINUS.text},operand={$f.st})
       | p=power {$type = $p.type;} -> {$p.st}
       ;

power returns [Symbol type]
      : atom (options {greedy=true;}: DSTAR f=factor)? {if(f!=null){$type = resolveType("power", $atom.type, $factor.type);} else {$type = $atom.type;}} 
        -> power(atom={$atom.st}, pow={$DSTAR.text}, factor={$f.st}) //if both int types, cast to int
      ; 
      
subscript : LBRACK expr RBRACK -> template(sub={$expr.st}) "[<sub>]"
          ;
              
     
atom returns [Symbol type]
@init{
  FunctionSymbol func = null;
}
      : f=NAME LPAREN (al=arglist)? RPAREN {Symbol s = handleFunctionCall($f.text, $f.line, $al.args); if(s == null){doError($f.line, "function '"+ $f.text +"' not defined."); $type = null;} else $type=s.type;}
      {func = (FunctionSymbol) currentScope.lookup($f.text);}
      -> function_call(name={$f.text}, args={$al.st}, bool={func.isRecursive()})
//     |  a=NAME (su+=subscript)+ {Symbol s = currentScope.lookup($a.text); if (s == null) System.out.println("ERROR! list '" + $a.text + "' not defined"); $type = s.type;}
//      -> array_ref(name={$a.text}, subs={$su}) 
     |  v=NAME {Symbol s = currentScope.lookup($v.text); if (s==null) doError($v.line, "variable '"+ $v.text + "' not defined."); $type = s.type;} 
      -> {new StringTemplate("fire_"+$NAME.getText())}
     | INT {$type = IntTypeSymbol.instance();} -> {new StringTemplate($INT.text)}
     | FLOAT {$type = FloatTypeSymbol.instance();} -> {new StringTemplate($FLOAT.text)}
     | TRUE {$type = BoolTypeSymbol.instance();} -> {new StringTemplate($TRUE.text)}
     | FALSE {$type = BoolTypeSymbol.instance();} -> {new StringTemplate($FALSE.text)}
     | NULL {$type = NullTypeSymbol.instance();} -> {new StringTemplate($NULL.text)}
     | LPAREN t=test RPAREN {$type = $test.type;} -> template(test={$t.st}) "(<test>)"
     | STRING {$type = StringTypeSymbol.instance();} -> {new StringTemplate($STRING.text)}
     ;

arglist returns [ArrayList<Symbol> args] 
@init{
  $args = new ArrayList<Symbol>();
   ArrayList<StringTemplate> t2 = new ArrayList<StringTemplate>();
}
         : t=test {$args.add($t.type); t2.add($t.st);} (o+=COMMA t1=test {$args.add($t1.type); t2.add($t1.st);})* ->argument_list(list={t2}) 
         ;
         
parameters returns [ArrayList<Symbol> args] 
@init{ 
  $args = new ArrayList<Symbol>();
  ArrayList<StringTemplate> l2 = new ArrayList<StringTemplate>();
 } 
            : (l=defparameter {$args.add($l.type); l2.add($l.st); if(inRecursiveFunc) paramList.add($l.type.name);})? 
              (COMMA l1=defparameter {$args.add($l1.type); l2.add($l1.st); if(inRecursiveFunc) paramList.add($l1.type.name);})*
              -> paramlist(list={l2})
            ;
           
defparameter returns [Symbol type]  
      : t=type NAME {$type = t.type; currentScope.defineVariable($NAME.text,t.type);} 
                -> defparam(type={t.st}, var={$NAME.text}) 
      ;

type returns [Symbol type] 
      : c=ctype LT t=type GT { if($c.tname.equals("list")){$type = ListTypeSymbol.instance($t.type);}}
        -> type(container={$c.st}, type={$t.st})
      | p=ptype {$type = $p.type;}-> {$p.st}
      ;
     
ctype returns [String tname] : TLIST {$tname = new String("list");} -> tlist()
                             ;


ptype returns [Symbol type]
      : TINT {$type = IntTypeSymbol.instance();} ->  tint()
      | TFLOAT {$type = FloatTypeSymbol.instance();} ->  tfloat() 
      | TBOOL {$type = BoolTypeSymbol.instance();} ->  tbool()
      ;
     
     
/*keywords*/
IF  : 'if';
ELSE  : 'else';
ELSEIF  : 'elif';
WHILE : 'while';
FOR : 'for';
IN : 'in';
RANGE : 'range';
PRINT : 'print';
DEFINE : 'define';
RETURN  : 'return';
RECURSIVE : 'recursive';
BASE : 'base';
WATCH  : 'watch';
     
 /*TYPES*/  
TINT : 'int' ;
TFLOAT : 'float';
TLIST : 'list';
TBOOL : 'bool';

/*keywords*/
TRUE : 'true';
FALSE : 'false';
NULL : 'null';

/* grouping */
LPAREN  : '(';
RPAREN  : ')';
LBRACK  : '[';
RBRACK  : ']';
LCURLY  : '{';
RCURLY  : '}';

/* punctuation */
COLON : ':';
COMMA : ',';
DOT : '.';

/* comparison */
LT  : '<';
GT  : '>';
EQ  : '==';
LEQ : '<=';
GEQ : '>=';
NEQ : '!=';

/* short circuit */
AND : 'and';
OR  : 'or';
NOT : 'not';


/* operators */
PLUS  : '+';
MINUS : '-';
STAR  : '*';
SLASH : '/';
PERCENT : '%';
DSTAR : '**';


/* assignment */
ASSIGN  : '=';
PLUSEQ  : '+=';
MINUSEQ : '-=';
STAREQ  : '*=';
SLASHEQ : '/=';


/* types */

fragment
EXPONENT
    :    ('e' | 'E') ( '+' | '-' )? DIGITS
    ;

fragment
DIGITS  : ('0'..'9')+;
    
FLOAT
    :   '.' DIGITS (EXPONENT)?
    |   DIGITS '.' EXPONENT
    |   DIGITS ('.' (DIGITS (EXPONENT)?)? | EXPONENT)
    ;
    
INT :  '0' ('x' | 'X') ( '0' .. '9' | 'a' .. 'f' | 'A' .. 'F' )+ // Hex
    |  '0' DIGITS*  // Octal and Regular Zero
    |  '1'..'9' DIGITS* // Regular
    ;

NAME:    ( 'a' .. 'z' | 'A' .. 'Z' | '_')
        ( 'a' .. 'z' | 'A' .. 'Z' | '_' | '0' .. '9' )* 
    ;
    
STRING  : '"' (ESC | ~('\\'|'\n'|'"'))* '"' ;


ESC : '\\' . ;

CONTINUED_LINE
    :    '\\' (' '|'\t')* ('\r')? '\n' (' '|'\t')*  { $channel=HIDDEN; }
         ( nl=NEWLINE {emit(new ClassicToken(NEWLINE,nl.getText()));} )?
    ;
    
COMMENT
@init{
    _channel=HIDDEN;
}
    :    {getCharPositionInLine()==0}?=> (' '|'\t')* '#' (~'\n')* ('\n')* 
    |    {getCharPositionInLine()>0}?=> '#' (~'\n')* // let NEWLINE handle \n unless char pos==0 for '#'
    ;

NEWLINE
    :   {getCharPositionInLine()==0}?=> (' ' | '\t')* (('\r')? '\n' )+ {$channel=HIDDEN;}
    |   {getCharPositionInLine()>0}?=> (('\r')? '\n' )+
    ;

WS  :    {getCharPositionInLine()>0}?=> (' '|'\t')+ {$channel=HIDDEN;}
    ;

LEADING_WS
@init {
    int spaces = 0;
}
    :   {getCharPositionInLine()==0}?=>
        ( (' '  { spaces++; } 
          |'\t' { spaces+= TAB_WIDTH; spaces -= (spaces \% TAB_WIDTH); }
          )+
          {
            char[] ind = new char[spaces];
            for (int i = 0; i<spaces; i++) ind[i] = ' ';
            emit(new ClassicToken(LEADING_WS, new String(ind)));
          }
          ( ('\r')? '\n' {if (state.token!=null) state.token.setChannel(HIDDEN); else $channel=HIDDEN;})*
        )
    ;
    
